import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.collections import PatchCollection
import matplotlib.animation as animation
import sympy as sy
import simtk.unit as unit
from simtk import openmm as mm
from simtk.openmm import app
import skopt as skopt
from tqdm import tqdm

2D Lennard-Jones fluid

mass = 39.948 * unit.amu
sigma = 3.404 * unit.angstroms
epsilon = 0.238 * unit.kilocalories_per_mole
charge = 0.0 * unit.elementary_charge

n_particles = 50
reduced_density = 0.85
l_box = None # 40.0 * unit.angstroms

temperature = 300.00 * unit.kelvin
integration_timestep = 0.002 * unit.picoseconds
collisions_rate = 1.0 / unit.picoseconds

production_time = 0.1 * unit.nanoseconds
saving_time = 0.25 * unit.picoseconds

radius = 2**(-5.0/6.0)*sigma

if l_box is None:
    area_particles = n_particles*np.pi*radius**2
    area = area_particles/reduced_density
    l_box = area**(1/2)
    print('Side of the box: {}'.format(l_box))
else:
    area_particles = n_particles*np.pi*radius**2
    area = l_box**2
    reduced_density = area_particles/area
    print('Reduced density: {}'.format(reduced_density))

Side of the box: 25.97058290208812 A

space = skopt.Space([[0.0, l_box._value], [0.0, l_box._value]])
generator = skopt.sampler.Lhs(criterion="maximin", iterations=10000)
positions_2d = generator.generate(space.dimensions, n_particles)
positions_2d = np.array(positions_2d)*unit.angstroms

system = mm.System()

for _ in range(n_particles):
    system.addParticle(mass)

v1 = np.zeros(3) * unit.angstroms
v2 = np.zeros(3) * unit.angstroms
v3 = np.zeros(3) * unit.angstroms

v1[0] = l_box
v2[1] = l_box
v3[2] = l_box

system.setDefaultPeriodicBoxVectors(v1, v2, v3)

non_bonded_force = mm.NonbondedForce()
non_bonded_force.setNonbondedMethod(mm.NonbondedForce.CutoffPeriodic)
non_bonded_force.setCutoffDistance(3.0*sigma)
non_bonded_force.setUseSwitchingFunction(True)
non_bonded_force.setSwitchingDistance(2.0*sigma)
non_bonded_force.setUseDispersionCorrection(True)

for _ in range(n_particles):
    non_bonded_force.addParticle(charge, sigma, epsilon)

_ = system.addForce(non_bonded_force)

armonic_force = mm.CustomExternalForce('A*periodicdistance(0, 0, z, 0, 0, 0)^2')
k=2500.0*unit.kilocalories_per_mole/unit.nanometers**2
armonic_force.addGlobalParameter('A', 0.5*k)
for ii in range(n_particles):
    armonic_force.addParticle(ii, [])
_ = system.addForce(armonic_force)

integrator = mm.LangevinIntegrator(temperature, collisions_rate, integration_timestep)
platform = mm.Platform.getPlatformByName('CUDA')
context = mm.Context(system, integrator, platform)

initial_positions=np.zeros([n_particles,3])*unit.angstroms
initial_positions[:,0:2]=positions_2d[:,:]

initial_velocities=np.zeros([n_particles,3])*unit.angstroms/unit.picoseconds

context.setPositions(initial_positions)
context.setVelocities(initial_velocities)

state=context.getState(getEnergy=True)
print("Before minimization: {}".format(state.getPotentialEnergy()))
mm.LocalEnergyMinimizer_minimize(context)
state=context.getState(getEnergy=True, getPositions=True)
print("After minimization: {}".format(state.getPotentialEnergy()))

Before minimization: 10915837.8931304 kJ/mol
After minimization: -142.6552094439215 kJ/mol

context.setVelocities(initial_velocities)

positions_after_minimization = state.getPositions(asNumpy=True).in_units_of(unit.angstroms)

fig, axes = plt.subplots(1, 2)

axes[0].set_aspect('equal', 'box')

patches=[]
for ii in range(n_particles):
    patches.append(Circle(positions_2d[ii,:]._value, radius._value))

p = PatchCollection(patches, alpha=0.5)
axes[0].add_collection(p)

axes[0].set_title("Initial positions")
axes[0].set_xlim([0.0, l_box._value])
axes[0].set_ylim([0.0, l_box._value])

axes[1].set_aspect('equal', 'box')

patches=[]
for ii in range(n_particles):
    patches.append(Circle(positions_after_minimization[ii,0:2]._value, radius._value))

p = PatchCollection(patches, alpha=0.5)
axes[1].add_collection(p)

axes[1].set_title("Minimized positions")
axes[1].set_xlim([0.0, l_box._value])
axes[1].set_ylim([0.0, l_box._value])

plt.show()

production_n_steps = int(production_time/integration_timestep)
saving_n_steps = int(saving_time/integration_timestep)
n_saving_periods = int(production_n_steps/saving_n_steps)

time = np.zeros([n_saving_periods+1]) * unit.nanoseconds
trajectory = np.zeros([n_saving_periods+1, n_particles, 3]) * unit.angstroms
potential_energy = np.zeros([n_saving_periods+1]) * unit.kilocalories_per_mole
kinetic_energy = np.zeros([n_saving_periods+1]) * unit.kilocalories_per_mole

state = context.getState(getPositions=True, getEnergy=True, enforcePeriodicBox = True)
time[0] = state.getTime()
trajectory[0,:,:] = state.getPositions(asNumpy=True)
potential_energy[0] = state.getPotentialEnergy()
kinetic_energy[0] = state.getKineticEnergy()
for ii in tqdm(range(1,n_saving_periods+1)):
    integrator.step(saving_n_steps)
    state = context.getState(getPositions=True, getEnergy=True, enforcePeriodicBox = True)
    time[ii] = state.getTime()
    trajectory[ii,:,:] = state.getPositions(asNumpy=True)
    potential_energy[ii] = state.getPotentialEnergy()
    kinetic_energy[ii] = state.getKineticEnergy()

100%|██████████| 400/400 [00:02<00:00, 187.36it/s]

trajectory_mem = trajectory.size * trajectory.itemsize * unit.bytes
print('Trajectory size: {} MB'.format(trajectory_mem._value/(1024*1024)))

Trajectory size: 0.4589080810546875 MB

plt.rcParams["animation.html"] = "jshtml"
fig, ax = plt.subplots()

ax.set_aspect('equal', 'box')
ax.set_title("2D LJ Fluid")
ax.set_xlim([0.0, l_box._value])
ax.set_ylim([0.0, l_box._value])


frame=0
patches=[]
for ii in range(n_particles):
    patches.append(Circle(trajectory[frame,ii,0:2]._value, radius._value))

p = PatchCollection(patches, alpha=0.5)
ax.add_collection(p)

def animate(frame):

    global trajectory, radius, p

    patches=[]
    for ii in range(n_particles):
        patches.append(Circle(trajectory[frame,ii,0:2]._value, radius._value))

    p.set_paths(patches)

ani = animation.FuncAnimation(fig, animate, frames=trajectory.shape[0], interval=100, blit=False)
plt.close()
ani

Once Loop Reflect